package com.llj.chatbidemo.config;

import com.volcengine.ark.runtime.model.completion.chat.*;
import com.volcengine.ark.runtime.service.ArkService;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;
import javax.annotation.PostConstruct;
import java.util.ArrayList;
import java.util.List;

@Component
public class ChatAITool {

    @Value("${ai.volc.api-key}")
    private String apiKey;

    @Value("${ai.volc.base-url}")
    private String baseUrl;

    @Value("${ai.volc.model-id}")
    private String modelId;

    @Value("${ai.volc.system-prompt}")
    private String systemPrompt;

    private ArkService arkService;

    @PostConstruct
    private void init() {
        System.out.println(apiKey);
        this.arkService = ArkService.builder()
                .apiKey(apiKey)
                .baseUrl(baseUrl)
                .build();
    }

    /**
     * 标准问答（自动包含系统提示）
     */
    public String askQuestion(String userQuestion) {
        List<ChatMessage> messages = buildMessages(userQuestion);
        return executeChatCompletion(messages);
    }

    /**
     * 自定义消息列表问答
     */
    public String askWithMessages(List<ChatMessage> customMessages) {
        return executeChatCompletion(customMessages);
    }

    private List<ChatMessage> buildMessages(String userQuestion) {
        List<ChatMessage> messages = new ArrayList<>();
        messages.add(ChatMessage.builder()
                .role(ChatMessageRole.SYSTEM)
                .content(systemPrompt)
                .build());
        messages.add(ChatMessage.builder()
                .role(ChatMessageRole.USER)
                .content(userQuestion)
                .build());
        return messages;
    }

    private String executeChatCompletion(List<ChatMessage> messages) {
        try {
            ChatCompletionRequest request = ChatCompletionRequest.builder()
                    .model(modelId)
                    .messages(messages)
                    .build();

            return (String) arkService.createChatCompletion(request)
                    .getChoices()
                    .stream()
                    .findFirst()
                    .map(choice -> choice.getMessage().getContent())
                    .orElseThrow(() -> new RuntimeException("AI响应为空"));
        } catch (Exception e) {
            throw new ChatException("AI服务调用失败", e);
        }
    }

    public void shutdown() {
        if (arkService != null) {
            arkService.shutdownExecutor();
        }
    }

    // 自定义异常类
    public static class ChatException extends RuntimeException {
        public ChatException(String message, Throwable cause) {
            super(message, cause);
        }
    }
}
